package com.example.shardingwythe.config;

import com.google.common.collect.Range;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.standard.RangeShardingValue;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.Collection;
import java.util.LinkedHashSet;

/**
 * @author 二师兄
 * @date 2021/9/7  11:49 上午
 */
@Component
@Slf4j
public class TableRangeShardingAlgorithm implements RangeShardingAlgorithm<LocalDateTime> {

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames,
                                         RangeShardingValue<LocalDateTime> rangeShardingValue) {
        Range<LocalDateTime> ranges = rangeShardingValue.getValueRange();

        LocalDateTime start = ranges.lowerEndpoint();
        LocalDateTime end = ranges.upperEndpoint();

        int startYear = start.getYear();
        int endYear = end.getYear();

        int startMonth = start.getMonthValue();
        int endMonth = end.getMonthValue();

        Collection<String> tables = new LinkedHashSet<>();
        if (start.getNano() <= end.getNano()) {
            for (String c : availableTargetNames) {
                int cMonth = Integer.parseInt(c.substring(c.length() - 6));
                if (cMonth >= Integer.parseInt("" + startYear + startMonth) && cMonth <= Integer.parseInt("" + endYear + endMonth)) {
                    tables.add(c);
                }
            }
        }
        return tables;
    }
}
